# -*- coding : utf-8 -*-
# @Time    : 2023/12/07 17:20
# @Author  : wanggang
# @Email   : wanggang2@lingxi.ai
# @File    : data_handling.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import os
import sys
from openpyxl import load_workbook, Workbook



class DataHandle(object):


    def __init__(self):
        self.update_data_from_file()


    def update_data_from_file(self):

        customer_info_dict = dict()
        for file_name in os.listdir('./data/'):
            if file_name.endswith('xlsx'):
                file_path = os.path.join('./data/', file_name)
                wb = load_workbook(file_path)
                ws = wb[wb.sheetnames[0]]

                for i, row in enumerate(ws.values):
                    if i != 0:
                        customer_id = row[0]
                        call_start_time = row[1]
                        session_id = row[2]
                        dialogue = row[4]
                        base_info = row[5]

                        dialogue_result_list = list()
                        dialogue_list = dialogue.split('\n')
                        for dialogue in dialogue_list:
                            if dialogue.startswith('销售员'):
                                dialogue_result_list.append({
                                    'role': 'assistant',
                                    'content': dialogue.replace('销售员:')
                                })
                            elif dialogue.startswith('用户'):
                                dialogue_result_list.append({
                                    'role': 'user',
                                    'content': dialogue.replace('用户:')
                                })

                        customer_info_dict[customer_id] = {
                            'call_start_time': call_start_time,
                            'session_id': session_id,
                            'dialogue': dialogue_result_list,
                            'base_info': base_info
                        }

        self.customer_info_dict = customer_info_dict


    def get_customer_info(self, customer_id):

        return self.customer_info_dict[customer_id]


if __name__ == '__main__':


    for file_name in os.listdir('../data/'):
        print(os.path.join('../data/', file_name))

    print( os.listdir('../data/') )
    print(os.path.dirname(os.path.realpath(__file__)))


